# -*- coding: utf-8 -*-
"""
@Time ： 2025/7/26 0:07
@Auth ： 韩明航
"""
from langchain.docstore.document import Document
from langchain.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import numpy as np

embedding_model = HuggingFaceEmbeddings(model_name="chuxin-llm/Chuxin-Embedding")

vectorstore = FAISS.load_local("faiss_index", embedding_model, allow_dangerous_deserialization=True)

retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

query = "哪些景点热度比较高，并介绍一下"

# 检索与 query 相关的文档片段
retrieved_docs = retriever.invoke(query)

# 打印检索到的文档片段
for i, doc in enumerate(retrieved_docs):
    print(f"Document {i + 1}:")
    print(f"Content: {doc.page_content}\n")

from transformers import AutoTokenizer, AutoModelForCausalLM


# 远程
model_path = 'neuralmagic/Mistral-7B-Instruct-v0.3-GPTQ-4bit'


# 加载
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype="auto",  # 自动选择模型的权重数据类型
    device_map="auto",   # 自动选择可用的设备（CPU/GPU）
)

from transformers import pipeline

generator = pipeline(
    "text-generation",  # 指定任务类型为文本生成
    model=model,
    tokenizer=tokenizer,
    max_length=4096,    # 指定生成文本的最大长度
    pad_token_id=tokenizer.eos_token_id
)

from langchain_huggingface import HuggingFacePipeline

llm = HuggingFacePipeline(pipeline=generator)

from langchain.prompts import PromptTemplate

custom_prompt = PromptTemplate(
    template="""请你根据我的要求推荐关于四川的景点。以下是关于四川景点的信息：
{context}

根据上述信息，请回答以下问题：
{question}""",
    input_variables=["context", "question"]
)

from langchain.chains import RetrievalQA

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",   # 直接堆叠所有检索到的文档
    retriever=retriever,  # 使用先前定义的检索器来获取相关文档

    chain_type_kwargs={"prompt": custom_prompt}
)

# 提出问题
query = "给我推荐热度较高的景点"

# 获取答
try:
    # 在调用模型生成函数时设置 max_length
    answer = qa_chain.invoke(query)
    print("生成的回答：", answer['result'])
except Exception as e:
    print("发生错误：", str(e))